# -*- coding: utf-8 -*-
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""
变量声明分析 Pass

在代码生成之前分析哪些赋值语句是变量声明，避免在代码生成阶段重复分析。
处理复杂的变量声明判断逻辑，包括导入变量、global 语句、多目标赋值等。
"""

from typing import Dict, Set, Optional, Any, List, Union
from .base import BasePass, PassStage
from ..ir.base import IRNode
from ..ir.module import IRModule
from ..ir.statements.assignment import IRAssign
from ..ir.statements.global_stmt import IRGlobal, IRNonlocal
from ..ir.declarations.function import IRFunctionNode
from ..ir.expressions.basic import IRName
from ..symbol_table import SymbolKind


class DeclarationAnalysisPass(BasePass):
    """变量声明分析 Pass"""

    name = "variable_declaration_analysis"
    stage = PassStage.POST_ANALYSIS

    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__()
        self.config = config or {}
        self.scope_stack: List[Dict[str, Any]] = []
        self.current_scope_node: Optional[IRNode] = None  # 当前正在分析的作用域节点

    def run(self, ir_module: IRModule, context) -> Any:
        """运行变量声明分析 Pass"""
        # 初始化作用域栈
        self.scope_stack = []

        # 从根节点开始递归分析
        self._analyze_scope(ir_module, ir_module)

        return ir_module

    def _analyze_scope(self, scope_node: IRNode, ir_module: IRModule):
        """分析一个作用域节点

        Args:
            scope_node: 作用域节点（模块、函数、类等）
        """
        # 保存之前的作用域节点
        previous_scope_node = self.current_scope_node
        self.current_scope_node = scope_node

        # 进入新作用域
        scope_info = {
            "declared_vars": {},  # 当前作用域已声明的变量: {var_name: IRNode}
            "global_vars": set(),  # 当前作用域的 global 声明
            "scope_type": scope_node.node_type,  # 作用域类型
            "scope_node": scope_node,  # 作用域节点引用
        }

        # 根据作用域类型初始化
        if isinstance(scope_node, IRFunctionNode):
            for param in scope_node.params:
                # 函数参数一定是变量声明，先搜集起来
                scope_info["declared_vars"][param.name] = scope_node
            # 全局变量声明也先搜集起来
            global_vars = self._get_global_declarations(scope_node, ir_module)
            scope_info["global_vars"] = global_vars
            # 将全局变量信息存储到函数节点的扩展参数中，供后续分析使用
            scope_node.set_extension_param("global_vars", global_vars)

        elif scope_node.node_type == "class":
            # 类作用域：类变量在类体中声明
            pass

        elif scope_node.node_type == "module":
            # 模块作用域：全局作用域
            pass

        elif scope_node.node_type in (
            "list_comp",
            "set_comp",
            "dict_comp",
            "generator_exp",
            "comprehension",
        ):
            # 推导式作用域：有自己的局部变量
            pass

        self.scope_stack.append(scope_info)

        try:
            self._analyze_in_current_scope(scope_node, ir_module)
        finally:
            # 退出作用域
            self.scope_stack.pop()
            # 恢复之前的作用域节点
            self.current_scope_node = previous_scope_node

    def _analyze_in_current_scope(self, node: IRNode, ir_module: IRModule):
        """在当前作用域中继续分析节点"""
        for child in node.get_child_nodes():
            if isinstance(child, IRAssign):
                self._analyze_assignment(child, ir_module)
            elif isinstance(child, IRGlobal):
                # 处理 global 语句
                self._analyze_global_statement(child, ir_module)
            elif isinstance(child, IRNonlocal):
                # 处理 nonlocal 语句
                self._analyze_nonlocal_statement(child, ir_module)
            elif child.creates_scope():
                self._analyze_scope(child, ir_module)
            else:
                # 检查变量访问
                self._check_variable_access(child, ir_module)
                self._analyze_in_current_scope(child, ir_module)

    def _analyze_assignment(self, assign_node: IRAssign, ir_module: IRModule):
        """分析赋值语句

        Args:
            assign_node: 赋值节点
        """
        if not self.scope_stack:
            return

        current_scope = self.scope_stack[-1]
        targets = self._extract_assignment_targets(assign_node)

        for target in targets:
            # 检查是否是变量声明还是变量赋值
            is_declaration = self._is_variable_declaration(
                target, assign_node, ir_module
            )

            # 设置扩展参数
            if len(targets) > 1:
                # 多目标赋值
                target_declarations = assign_node.get_extension_param(
                    "target_declarations", {}
                )
                target_declarations[target] = is_declaration
                assign_node.set_extension_param(
                    "target_declarations", target_declarations
                )
            else:
                # 单目标赋值
                assign_node.set_extension_param(
                    "is_variable_declaration", is_declaration
                )

            # 如果是声明，添加到当前作用域
            if is_declaration:
                current_scope["declared_vars"][target] = assign_node

    def _is_variable_declaration(
        self, target: str, assign_node: IRAssign, ir_module: IRModule
    ) -> bool:
        """判断是否是变量声明

        Args:
            target: 目标变量名
            assign_node: 赋值节点

        Returns:
            bool: 是否是变量声明
        """
        if not self.scope_stack:
            return True

        current_scope = self.scope_stack[-1]
        scope_type = current_scope["scope_type"]

        # 1. 检查是否是 global 变量
        # 对于非函数作用域（如 exception_handler），需要向上查找到函数作用域的 global 声明
        if self._is_global_variable(target):
            return False  # global 变量的赋值不是声明

        # 2. 检查是否已在当前作用域声明
        if target in current_scope["declared_vars"]:
            # 已经在当前作用域声明，是赋值不是声明
            return False

        # 3. 检查是否是导入的符号
        if self._is_imported_symbol(target, assign_node):
            # 外部导入的符号只能是赋值
            return False

            # 4. 类作用域的特殊处理
        if scope_type == "class":
            # 在类作用域中，赋值通常是类变量声明
            # 除非是在方法中（但那会在方法作用域中处理）
            return True

        # 5. 检查是否是外层作用域的变量（闭包变量）
        if self._is_closure_variable(target):
            # 在推导式中，外层变量可能被重新绑定
            if scope_type in (
                "list_comp",
                "set_comp",
                "dict_comp",
                "generator_exp",
                "comprehension",
            ):
                return True  # 推导式中的变量是局部的
            return False

        # 6. 其他情况都是新变量声明
        return True

    def _is_global_variable(self, target: str) -> bool:
        """检查变量是否在某个父函数作用域中声明为 global

        向上遍历作用域栈，查找包含该变量 global 声明的函数作用域。
        这样可以正确处理 exception_handler 等嵌套作用域中的 global 变量。

        Args:
            target: 变量名

        Returns:
            bool: 是否是 global 变量
        """
        # 从当前作用域向上查找
        for scope in reversed(self.scope_stack):
            # 检查当前作用域是否有 global 声明
            if target in scope["global_vars"]:
                return True

            # 如果遇到函数作用域但没有找到 global 声明，说明不是 global
            # （因为 global 声明只在声明它的函数作用域有效）
            if scope["scope_type"] == "function":
                return False

        return False

    def _is_closure_variable(self, target: str) -> bool:
        """检查是否是闭包变量（外层作用域的变量）

        Args:
            target: 变量名

        Returns:
            bool: 是否是闭包变量
        """
        # 从外层作用域向上查找
        for scope in reversed(self.scope_stack[:-1]):  # 排除当前作用域
            if target in scope["declared_vars"]:
                return True
        return False

    def _get_global_declarations(
        self, scope_node: IRNode, ir_module: IRModule
    ) -> Set[str]:
        """获取作用域的 global 声明

        Args:
            scope_node: 作用域节点（函数、方法等）

        Returns:
            Set[str]: global 变量集合
        """
        global_vars = set()

        # 只有函数和方法作用域才有 global 声明
        if scope_node.node_type in ("function", "method"):
            self._get_global_statements(scope_node, global_vars, ir_module)

        return global_vars

    def _get_global_statements(
        self, node: IRNode, global_vars: set, ir_module: IRModule
    ):
        """在节点中查找 global 语句

        Args:
            node: 要搜索的节点

        Returns:
            Set[str]: 找到的 global 变量集合
        """
        for child in node.get_child_nodes():
            if isinstance(child, IRGlobal):
                global_vars.update(child.names)
            # 递归查找，但不跨越作用域
            elif not child.creates_scope():
                self._get_global_statements(child, global_vars, ir_module)

    def _ensure_global_node(self, node: IRGlobal, ir_module: IRModule):
        """验证 global 节点是否存在

        Args:
            node: global 节点
            ir_module: 模块
        """
        for name in node.names:
            if not self._validate_global_variable(name, ir_module):
                raise ValueError(
                    f"模块 {ir_module.module_full_name} 中 global 变量 {name} 不存在"
                )

    def _validate_global_variable(self, var_name: str, ir_module: IRModule) -> bool:
        """验证 global 变量在外部作用域是否存在

        Args:
            var_name: 变量名
            assign_node: 赋值节点

        Returns:
            bool: 是否存在
        """
        # 只调用基类的函数查找模块本地符号，不递归到上层模块
        symbol = IRNode.lookup_symbol(ir_module, var_name)
        return symbol is not None and symbol.kind == SymbolKind.VARIABLE

    def _extract_assignment_targets(self, assign_node: IRAssign) -> List[str]:
        """提取赋值语句的所有目标

        Args:
            assign_node: 赋值节点

        Returns:
            List[str]: 目标变量名列表
        """

        targets = []
        target_exprs = assign_node.get_target_exprs()
        for expr in target_exprs:
            # 除了 IRName 外，其他赋值表达式都不可能是变量声明
            if isinstance(expr, IRName) and expr.id.isidentifier():
                targets.append(expr.id)
        return targets

    def _analyze_global_statement(self, global_node: IRGlobal, ir_module: IRModule):
        """分析 global 语句

        Args:
            global_node: global 语句节点
        """
        if not self.scope_stack:
            return

        current_scope = self.scope_stack[-1]
        # 验证 global 变量是否在模块级别存在
        for name in global_node.names:
            # 先检查是否已经在当前作用域的 global_vars 中（提前收集的）
            if name not in current_scope["global_vars"]:
                # 如果不在，验证是否在模块级别存在
                if not self._validate_global_variable(name, ir_module):
                    raise ValueError(
                        f"模块 {ir_module.module_full_name} 中 global 变量 {name} 不存在"
                    )

    def _analyze_nonlocal_statement(
        self, nonlocal_node: IRNonlocal, ir_module: IRModule
    ):
        """分析 nonlocal 语句

        Args:
            nonlocal_node: nonlocal 语句节点
        """
        if not self.scope_stack:
            return

        # nonlocal 语句的处理逻辑
        # 目前暂时跳过，因为 C++ 中没有直接对应的概念
        # TODO: 根据需要实现 nonlocal 的处理逻辑
        pass

    def _is_imported_symbol(self, name: str, assign_node: IRAssign) -> bool:
        """检查是否是导入的符号

        Args:
            name: 符号名
            assign_node: 赋值节点

        Returns:
            bool: 是否是导入的符号
        """
        # 查找符号
        symbol = assign_node.lookup_symbol(name)
        if not symbol or not symbol.ir_node:
            return False

        if symbol.ir_node.parent and symbol.ir_node.parent.node_type == "import":
            return True

        return False

    def _is_assignment_target(self, name_node: IRName) -> bool:
        """判断变量名是否是赋值目标（定义点）

        Args:
            name_node: 变量名节点

        Returns:
            bool: 如果是赋值目标返回 True
        """
        # 检查父节点是否是 IRAssign
        parent = name_node.parent
        if isinstance(parent, IRAssign):
            # 检查是否在 target_exprs 中
            for target in parent.target_exprs:
                if target is name_node:
                    return True
        return False

    def _check_variable_access(self, node: IRNode, ir_module: IRModule):
        """检查变量访问是否合法

        Args:
            node: 当前节点
            ir_module: IR 模块
        """
        # 检查是否是变量名节点
        if isinstance(node, IRName):
            # 跳过赋值语句的目标变量（它们是声明/定义，不是使用）
            if self._is_assignment_target(node):
                return
            var_name = node.id
            self._validate_variable_access(var_name, node, ir_module)

    def _validate_variable_access(
        self, var_name: str, access_node: IRNode, ir_module: IRModule
    ):
        """验证变量访问是否合法

        Args:
            var_name: 变量名
            access_node: 访问变量的节点
            ir_module: IR 模块
        """
        if not self.scope_stack:
            return

        current_scope = self.scope_stack[-1]

        # 1. 检查是否是 global 变量
        if var_name in current_scope["global_vars"]:
            return  # global 变量总是可以访问

        # 2. 只检查当前作用域内的变量声明
        # 因为上层作用域的变量都是合法的，所以只需要检查当前作用域
        declaration_node = current_scope["declared_vars"].get(var_name)

        if not declaration_node:
            return  # 变量不在当前作用域声明，可能在上层作用域、导入的或内置的，都不检查

        # 3. 检查当前作用域内声明节点是否在访问节点的父路径上
        if not self._is_declaration_accessible_in_scope(declaration_node, access_node):
            # 编译错误：变量在声明前被访问或作用域不可达
            raise ValueError(f"变量 '{var_name}' 在声明前被访问或作用域不可达。")

    def _is_declaration_accessible(
        self, declaration_node: IRNode, access_node: IRNode
    ) -> bool:
        """检查声明节点是否在访问节点的父路径上

        Args:
            declaration_node: 变量声明的节点
            access_node: 访问变量的节点

        Returns:
            bool: 如果声明节点在访问节点的父路径上返回True，否则返回False
        """
        # 从访问节点开始，沿着 parent 路径向上查找
        current = access_node.parent
        while current:
            if current == declaration_node:
                return True
            # 如果声明节点是赋值节点，检查其父节点
            if current == declaration_node.parent:
                return True
            current = current.parent

        return False

    def _is_declaration_accessible_in_scope(
        self, declaration_node: IRNode, access_node: IRNode
    ) -> bool:
        """检查声明节点是否在访问节点的父路径上，限制在当前作用域内

        遵循 Python 语义：
        - Python 的 if/else 不创建新作用域
        - 在 if/else 分支中声明的变量在函数作用域内都是可见的
        - 只要声明和访问都在同一个函数作用域内，就认为是可访问的

        Args:
            declaration_node: 变量声明的节点
            access_node: 访问变量的节点

        Returns:
            bool: 如果声明节点在访问节点的父路径上返回True，否则返回False
        """
        # 方案1：标准的父路径检查
        # 从访问节点开始，沿着 parent 路径向上查找，但不超过当前作用域节点
        current = access_node.parent
        while current and current != self.current_scope_node:
            if current == declaration_node:
                return True
            # 如果声明节点是赋值节点，检查其父节点
            if current == declaration_node.parent:
                return True
            current = current.parent

        # 方案2：检查是否到达了当前作用域节点
        # 如果到达了当前作用域，检查声明节点是否也在当前作用域内
        if current == self.current_scope_node:
            # 如果声明节点就是当前作用域节点或其直接子节点，也是可访问的
            if declaration_node == self.current_scope_node:
                return True
            if declaration_node.parent == self.current_scope_node:
                return True

            # 方案3：Python 语义 - if/else 不创建作用域
            # 检查声明节点是否也在当前函数作用域内（即使在不同的分支）
            # 这里简化检查：只要声明节点的某个祖先是当前作用域节点，就认为可访问
            decl_current = declaration_node.parent
            while decl_current:
                if decl_current == self.current_scope_node:
                    # 声明和访问都在同一个函数作用域内，Python 语义下是可访问的
                    return True
                decl_current = decl_current.parent

        return False
